#!/usr/bin/env python2.7
# Copyright (c) 2016 Presto Labs. Pte. Ltd.
# Author: jaewon

import subprocess
import sys
import tempfile


SECTION_BEGIN_PREFIX = "### BEGIN SECTION: "
SECTION_END_PREFIX = "### END SECTION: "


class SectionNotFound(Exception): pass
class UsageException(Exception): pass


def get_cron_config():
  conf_str = subprocess.check_output(['crontab', '-l'])
  return conf_str.split('\n')


def set_cron_config(cron_config):
  if get_cron_config() == cron_config:
    return
  else:
    f = tempfile.NamedTemporaryFile()
    try:
      f.write('\n'.join(cron_config))
      f.flush()
      subprocess.check_call(['crontab', f.name])
    finally:
      f.close()


def find_cron_section(section_id, cron_config=None):
  cron_config = cron_config or get_cron_config()
  begin_str = SECTION_BEGIN_PREFIX + str(section_id)
  end_str = SECTION_END_PREFIX + str(section_id)
  begin_lineno = None
  end_lineno = None
  for lineno, line in enumerate(cron_config):
    if line == begin_str:
      begin_lineno = lineno
    elif line == end_str:
      end_lineno = lineno
  if begin_lineno == None or end_lineno == None:
    raise SectionNotFound()
  return (begin_lineno, end_lineno)


def remove_cron_section(section_id, cron_config=None):
  cron_config = cron_config or get_cron_config()
  try:
    begin_lineno, end_lineno = find_cron_section(section_id, cron_config)
    return cron_config[0:begin_lineno] + cron_config[end_lineno+1:]
  except SectionNotFound:
    return cron_config


def add_cron_section(section_id, section_lines, cron_config=None):
  cron_config = cron_config or get_cron_config()
  cron_config = remove_cron_section(section_id, cron_config)
  return (cron_config
          + [SECTION_BEGIN_PREFIX + str(section_id)]
          + section_lines
          + [SECTION_END_PREFIX + str(section_id)])


def main(argv):
  try:
    if len(argv) != 3 and len(argv) != 4:
      raise UsageException()
    command = argv[1]
    section_id = argv[2].strip()

    if command == 'addsection':
      if len(argv) == 4:
        with open(argv[3], 'r') as f:
          section_lines = f.read().split('\n')
      else:
        # Reads from stdin.
        section_lines = sys.stdin.read().split('\n')
      if section_lines and section_lines[-1] == '':
        # Removes the last blank line.
        section_lines = section_lines[:-1]
      if not section_lines:
        return 0

      modified_cron = add_cron_section(section_id, section_lines)
      set_cron_config(modified_cron)
      return 0

    elif command == 'delsection' or command == 'rmsection':
      modified_cron = remove_cron_section(section_id)
      set_cron_config(modified_cron)
      return 0

    else:
      raise UsageException('unknown command: `%s`' % command)

  except UsageException as e:
    if e.message:
      sys.stderr.write(e.message + '\n')
    sys.stderr.write(
      ('Usage:\n'
       '  %s addsection <section-id> <filename>\n'
       '  %s delsection <section-id>\n') % (argv[0], argv[0]))
    return 1


if __name__ == '__main__':
  sys.exit(main(sys.argv))
